package org.alvin.code.jdbc.aop;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.Statement;

public class ConnectionProxyHandler implements InvocationHandler {
    private final Connection target;

    public ConnectionProxyHandler(Connection target) {
        this.target = target;
    }

    @Override
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
        if (method.getName().equals("createStatement")) {
            Statement statement = (Statement) method.invoke(target, args);
            return Proxy.newProxyInstance(
                    Statement.class.getClassLoader(),
                    new Class<?>[]{Statement.class},
                    new StatementProxyHandler(statement)
            );
        } else if (method.getName().equals("prepareStatement")) {
            PreparedStatement preparedStatement = (PreparedStatement) method.invoke(target, args);
            return Proxy.newProxyInstance(
                    PreparedStatement.class.getClassLoader(),
                    new Class<?>[]{PreparedStatement.class},
                    new PreparedStatementProxyHandler(preparedStatement)
            );
        }
        return method.invoke(target, args);
    }
}